import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import MDAnalysis as mda
import os
from Init import init_figs


def format_matrix_coord(x, y):
    col = int(x + 0.5)
    row = int(y + 0.5)
    if 0 <= col < numcols and 0 <= row < numrows:
        z = mat[row, col]
        return 'x=%1.4f, y=%1.4f, z=%1.4f' % (x, y, z)
    else:
        return 'x=%1.4f, y=%1.4f' % (x, y)


def draw_matrix_ax(ax, mat, extent, xlabel='Residue id', ylabel='Residue id'):
    h = ax.imshow(mat, extent=extent)
    numrows, numcols = mat.shape
    ax.format_coord = format_matrix_coord
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    plt.colorbar(h, ax=ax)
    return h


def draw_image(ax, img_path):
    ax.imshow(plt.imread(img_path))
    ax.set_xticks([])
    ax.set_yticks([])


if __name__ == '__main__':
    fig, axes = init_figs(1, 2)
    mat = pd.read_csv('test/matrix/residue matrix.csv', index_col=0)
    draw_matrix_ax(axes[0], mat, [10.5, 42.5, 42.5, 10.5])
    draw_matrix_ax(axes[1], mat, [10.5, 42.5, 42.5, 10.5])
    plt.tight_layout()
    plt.savefig('test_matrix.png')
